In [4]:
import cmocean as cmo
import numpy as np
import scipy.constants as spc
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.animation

from IPython.display import display, Math, HTML
from matplotlib.animation import FuncAnimation

plt.rcParams.update(
    {
        'mathtext.fontset':'cm', 
        'font.family':'serif', 
        'font.sans-serif':'Times New Roman', 
        'figure.figsize':[9,4], 
        'figure.titlesize':22,
        "figure.dpi":90, 
        'savefig.dpi':90,
        'axes.titlesize':20, 
        'axes.labelsize':18, 
        'axes.titley': 1.0, 
        'axes.titlepad': 5.0,
        'axes.edgecolor':'black', 
        'axes.grid': False, 
        'grid.alpha': .5,
        'xtick.labelsize':14,'ytick.labelsize':14,
        'xtick.major.size':6,'ytick.major.size':6,
        'xtick.major.width':1.25, 'ytick.major.width':1.25, 
        'xtick.direction':'inout','ytick.direction':'inout',
        'xtick.top':False, 'ytick.right':False,
        'legend.title_fontsize':14, 'legend.fontsize':14,
        'legend.borderaxespad': 1, 'legend.borderpad': 0.5,
        'legend.framealpha': 1,
        'legend.handleheight': 0.5, 'legend.handlelength': 2.0, 'legend.handletextpad':0.5,
        'legend.labelspacing': 0.25,
        'legend.fancybox':False,
        'legend.edgecolor': '0',
        'legend.frameon': True,
        'legend.markerscale': 1.25,
        'animation.embed_limit':2**128,
        'animation.html': 'jshtml'
    }
)

pi = spc.pi
%matplotlib inline
%config InlineBackend.figure_format = 'retina'

def animate2d(u, x, y, t, fps = 24, frames = 144):
    X, Y = np.meshgrid(x, y)
    t_steps = int(len(t)-1)
    i_frames = np.arange(0, t_steps, int(t_steps/frames))
    
    fig = plt.figure(figsize = (6,6), constrained_layout = True)
    ax = fig.add_subplot(111, projection = '3d')
    surf = ax.plot_surface(X, Y, u[0], alpha = 1) 
    def animate(i):
        ax.clear()
        surf = ax.plot_surface(X, Y, u[i], cmap = cmo.cm.solar, alpha = 1)
        ax.axis(False)
        ax.set_xlim(X.min(), X.max())
        ax.set_ylim(Y.min(), Y.max())
        ax.set_zlim(u[0].min(), u[0].max())   
        # ax.set_title(r'$t = {}$'.format(round(i*(t[1]-t[0]),1)))
        return surf, 
    anim = FuncAnimation(fig, animate, interval = int(1e3/fps), frames = i_frames, blit = False)
    plt.close()
    return anim

def periodic(f):
    f[0,:] = f[-2,:]
    f[-1,:] = f[2,:]
    f[:,0] = f[:,-2]
    f[:,-1] = f[:,2]
    return

def fixed(f):
    f[0,:] = 0
    f[-1,:] = 0
    f[:,0] = 0
    f[:,-1] = 0
    return

c = 1
x = np.linspace(0,1,128)
y = np.linspace(0,1,128)
xx, yy = np.meshgrid(x,y)

mask = ~((xx-.7)**2 + (yy-.5)**2 < .05)
u0 = 10/np.cosh(15*(xx-.3))**2 * 10/np.cosh(15*(yy-.5))**2
u0 = u0*mask

steps = 500

dx = x[1]-x[0]
dy = y[1]-y[0]
dt = (dx/c)/2
t = np.linspace(0,steps*dt,steps+1)
u = np.zeros(t.shape + u0.shape)
p = np.zeros(t.shape + u0.shape)
u[0] = u0
p[0] = np.zeros_like(u0)
rx = (c/dx)**2
ry = (c/dy)**2

boundary = 'fixed'

if boundary == 'periodic':
    for n in range(steps):
        u = u*mask
        periodic(u[n])
        p[n+1,1:-1,1:-1] = p[n,1:-1,1:-1] \
                    + dt*rx * (u[n,2:,1:-1] + u[n,:-2,1:-1] - 2*u[n,1:-1,1:-1]) \
                    + dt*ry * (u[n,1:-1,2:] + u[n,1:-1,:-2] - 2*u[n,1:-1,1:-1])
        u[n+1,1:-1,1:-1] = u[n,1:-1,1:-1] + dt*p[n+1,1:-1,1:-1]
        
if boundary == 'fixed':
    for n in range(steps):
        fixed(u[n])
        u = u*mask
        p[n+1,1:-1,1:-1] = p[n,1:-1,1:-1] \
                    + dt*rx * (u[n,2:,1:-1] + u[n,:-2,1:-1] - 2*u[n,1:-1,1:-1]) \
                    + dt*ry * (u[n,1:-1,2:] + u[n,1:-1,:-2] - 2*u[n,1:-1,1:-1])
        u[n+1,1:-1,1:-1] = u[n,1:-1,1:-1] + dt*p[n+1,1:-1,1:-1]
In [5]:
seconds = 5
fps = 24
frames = fps*seconds
anim = animate2d(u, x, y, t, fps = fps, frames = frames)
In [6]:
display(anim)
In [ ]: